# -*- encoding: utf-8 -*-

# @File    : generate_dataset_info.py
# @Time    : 19-6-20
# @Author  : zjh

r"""
Generate a dict format dataset info file contains `paths` `classes` and `display_names` of a dataset.
This scripts will find all images recursively under the give patch root directory with suffix of
"jpg", "jpeg", "png". And use each image's parent folder name as its class name. Then get its display_name
base on class name. the result file like:
```
{
    "paths":[".../ex_1.jpg",".../ex_1.png"],
    "classes":["foo", "bar"],
    "display_names":["biubiu", "xiuxiu"]
}
```
"""

import argparse
from collections import Counter

from .. import common


def get_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("patch_root",
                        help="The root directory of patches")
    parser.add_argument("output_path",
                        help="The output path of dataset info file.")
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Show details of dataset info.")
    return parser.parse_args()


FLAG = get_args()


def dataset_state(dataset):
    cnt = Counter(dataset["classes"])
    print("Example number: ", len(dataset["paths"]))
    print("Class number: ", len(cnt))
    print("Class distribution: ")
    fmt = "{:<8}{:<32}\t{:>8}\t{}"
    print(fmt.format("INDEX", "UUID", "COUNT", "NAME"))
    for i, (k, c) in enumerate(cnt.most_common()):
        print(fmt.format(i, k, c, k))


def main():
    img_ps = common.find_file_recursive(FLAG.patch_root, ["jpg", "jpeg", "png"],
                                        ignore_case=True)
    img_cs = [p.split("/")[-2] for p in img_ps]
    img_pys = [c for c in img_cs]

    dataset = {
        "paths": img_ps,
        "classes": img_cs,
        "display_names": img_pys
    }

    common.save_json(dataset, FLAG.output_path)

    if FLAG.verbose:
        dataset_state(dataset)


if __name__ == "__main__":
    main()
